from matplotlib import pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

from utils.get_stock_data import get_stock_data_table, gen_xls
import pandas as pd
import numpy as np


def machine_learn_trade(symbol, start_date, end_date):
    stock_dict = get_stock_data_table(symbol, start_date, end_date)
    stock = stock_dict['stock']
    print('股票代码：symbol', stock.head(3))
    x, y = classification_tc(stock)
    df = pd.concat([stock, x, y], axis=1)
    print('======df======')
    print(df.head(3))

    x_train, x_test, y_train, y_test = train_test_split(x, y, train_size=0.8)
    # 创建一个KNN实例
    knn_clf = KNeighborsClassifier(n_neighbors=95)
    knn_clf.fit(x_train, y_train)
    print('训练集中的准确率：', knn_clf.score(x_train, y_train))
    print('验证集中的准确率：', knn_clf.score(x_test, y_test))
    # 使用KNN模型预测每日股票的涨跌，保存为predict_signal
    # *** knn_clf.predict(x)内部的逻辑是什么，为什么返回的就是交易信号呢？
    df['predict_signal'] = knn_clf.predict(x)

    # *** 对数收益？这里书中没说明，对啥log(两个交易日相除)就是收益呢？
    df['return'] = np.log(df['close'] / df['close'].shift(1))  # 当日/前一天
    print('=============添加交易信号和收益后的结果====================')
    print(df.head(3))

    # 策略累计收益
    strategy_return_data = strategy_return(df, split_value=len(x_train))
    # 基础累计收益
    cum_return_data = cum_return(df, split_value=len(x_train))
    plot_chart(cum_return_data, strategy_return_data, 'dahua')


'''
    每日开盘价 - 每日收盘价， 并保存一个新的特征
    最高价 - 最低价，保存一个新的特征
    如果股票次日收盘价高于当日收盘价，标记为1， 代表次日股票价格上涨
    反之，如果次日收盘价格低于当日收半价格，标记为-1，代表次日股票价格下跌
    这个过程称之为：创建股票交易条件，trading condition
'''


# 定义一个用于分类的函数 classification_tc：预测股票下一个交易日是否上涨
# 给表增加三个字段
#     open - close 开盘价-收盘价
#     high - low 最高价-最低价
# 分类函数
def classification_tc(df):
    df['open-close'] = df['open'] - df['close']  # 开盘价-收盘价
    df['high-low'] = df['high'] = df['low']  # 最高价-最低价

    # 增加一个target字段，如果次日收盘价高于当日收盘价格，标记为1，否则为-1
    # df.shift(-1) 将列向上移动一位，这里就把后一天的和当天的数据对齐，然后直接做比较运算
    df['target'] = np.where(df['close'].shift(-1) > df['close'], 1, -1)
    # 去掉有空值的行
    df = df.dropna()
    # 将open-close和high-low作为数据集的特征
    x = df[['open-close', 'high-low']]
    y = df['target']
    return x, y


# 定义一个用于回归的函数:预测次日收盘价格与当日收盘价格之差
# 特征的添加与分类函数类似
# 只不过target字段改为 次日收盘价 - 当前收盘价 ---> df['close'].shift(-1) - df['close']
def regression_tc(df):
    df['open-close'] = df['open'] - df['close']  # 开盘价-收盘价
    df['high-low'] = df['high'] = df['low']  # 最高价-最低价
    df['target'] = df['close'].shift(-1) - df['close']
    df.dropna()
    # 将open-close和high-low作为数据集的特征
    x = df[['open-close', 'high-low']]
    y = df['target']
    return x, y


# 基础收益
def cum_return(df, split_value):
    print('cum_return split_value:', split_value)
    cum_return_data = df[split_value:]['return'].cumsum() * 100
    return cum_return_data


# 使用策略的收益
def strategy_return(df, split_value):
    print('strategy_return split_value:', split_value)
    df['strategy_return'] = df['return'] * df['predict_signal'].shift(1)
    cum_strategy_return = df[split_value:]['strategy_return'].cumsum() * 100

    return cum_strategy_return

# 绘制曲线
def plot_chart(cum_return_data, cum_strategy_return, symbol):
    plt.figure(figsize=(9, 6))
    plt.plot(cum_return_data, '--', label='%s 收益是' % symbol)

    plt.plot(cum_strategy_return, label='策略收益')

    plt.legend()
    plt.show()


if __name__ == '__main__':
    machine_learn_trade('002419', '20210101', '20221231')
